# -*- coding: utf-8 -*-

"""
@Datetime: 2019/3/30
@Author: Zhang Yafei
"""
import numpy as np
import pandas as pd
import scipy.stats
from scipy.stats import norm

p = np.asarray([0.65, 0.25, 0.07, 0.03])
q = np.array([0.6, 0.25, 0.1, 0.05])
q2 = np.array([0.1, 0.2, 0.3, 0.4])


def KL_divergence(p, q):
    """
    有时也称为相对熵，KL距离。对于两个概率分布P、Q，二者越相似，KL散度越大。
    KL散度满足非负性
    KL散度是不对称的，交换P、Q的位置将得到不同结果。
    :param p:
    :param q:
    :return:
    """
    return scipy.stats.entropy(p, q)


def JS_divergence(p, q):
    """
    JS散度基于KL散度，同样是二者越相似，JS散度越大。
        JS散度的取值范围在0-1之间，完全相同时为0
        JS散度是对称的
    :param p:
    :param q:
    :return:
    """
    M = (p + q) / 2
    return 0.5 * scipy.stats.entropy(p, M) + 0.5 * scipy.stats.entropy(q, M)


# 1000个均值170，标准差10的正态分布身高样本
h_real = norm.rvs(loc=170, scale=10, size=1000)
h_predict1 = norm.rvs(loc=168, scale=9, size=1000)
h_predict2 = norm.rvs(loc=160, scale=20, size=1000)


def JS_div(arr1, arr2, num_bins):
    max0 = max(np.max(arr1), np.max(arr2))
    min0 = min(np.min(arr1), np.min(arr2))
    bins = np.linspace(min0 - 1e-4, max0 - 1e-4, num=num_bins)
    PDF1 = pd.cut(arr1, bins).value_counts() / len(arr1)
    PDF2 = pd.cut(arr2, bins).value_counts() / len(arr2)
    return JS_divergence(PDF1.values, PDF2.values)


if __name__ == '__main__':
    print(KL_divergence(p, q))  # 0.011735745199107783
    print(KL_divergence(q, p))  # 0.013183150978050884
    print(JS_divergence(p, q))  # 0.003093977084273652
    print(JS_divergence(p, q2))  # 0.24719159952098618
    print(JS_divergence(p, p))  # 0.0
    print(JS_div(h_real, h_predict1, num_bins=20))  # 0.0098
    print(JS_div(h_real, h_predict2, num_bins=20))  # 0.135
